import json
import os
import sys
from typing import List
from importlib.metadata import distribution


def get_package_path(package_name):
    try:
        dist = distribution(package_name)
        return dist.locate_file("")
    except:
        return None


ascend_home_path = os.environ["ASCEND_HOME_PATH"]
file_path = os.path.dirname(__file__)
includes = [
    # 这三个是同一个目录
    f"{ascend_home_path}/include",
    f"{ascend_home_path}/runtime/include",
    f"{ascend_home_path}/acllib/include",
    f"{file_path}/../include",
    f"build_out/autogen",
]

include_cpu = [
    f"{ascend_home_path}/tools/tikicpulib/lib/include",
]

includes_compiler = [
    f"{ascend_home_path}/compiler/tikcpp/tikcfw/",
    f"{ascend_home_path}/compiler/tikcpp/tikcfw/impl",
    f"{ascend_home_path}/compiler/tikcpp/tikcfw/interface",
]


soc_flags = {
    "Ascend310B": "-D__DAV_M300__ -D__CCE_AICORE__=300",
    "Ascend910B": "-D__DAV_C220__ -D__CCE_AICORE__=220",
}


def get_pybind11_include():
    ret: str = os.popen("python3 -m pybind11 --includes").read()
    ret = ret.replace("\n", "")
    return ret


def get_torch_include():
    torch_path = get_package_path("torch")
    print("torch_path", torch_path)
    torch_npu_path = get_package_path("torch_npu")
    print("torch_npu_path", torch_npu_path)
    return [
        f"{torch_path}",
        f"{torch_npu_path}",
        f"{torch_path}/torch",
        f"{torch_path}/torch/include",
        f"{torch_path}/torch/csrc/api/include",
        f"{torch_npu_path}/torch_npu",
        f"{torch_npu_path}/torch_npu/include",
    ]


def gen_include_str(includes: List[str]):
    ret = ""
    for i in includes:
        ret += f" -I{i} "
    return ret


def get_cpps(path: str):
    f: str = ""
    ret = []
    for f in os.listdir(path):
        if f.endswith(".cpp"):
            ret.append(f"{path}/{f}")
    return ret


def gen_compile_commands(directory: str, cpps: List[str], soc_version):
    ret = []
    torch_include_paths = get_torch_include()
    include_str = gen_include_str(
        torch_include_paths + includes + include_cpu + includes_compiler
    )
    include_str = get_pybind11_include() + " " + include_str
    soc_flag = soc_flags.get(soc_version)
    if soc_flag is None:
        print(soc_version, "没找到")
    commands = f"/usr/bin/c++  -DASCENDC_CPU_DEBUG=1 -D__CCE_KT_TEST__=1 -std=c++17 -std=gnu++17 {soc_flag} {include_str}"
    for cpp in cpps:
        ret.append(
            {"directory": directory, "command": f"{commands} {cpp}", "file": cpp}
        )
    return ret


if __name__ == "__main__":
    directory = sys.argv[1]
    target_path = sys.argv[2]
    try:
        soc_version = sys.argv[3]
    except:
        soc_version = "Ascend310B"
    cpps = get_cpps(directory)
    commands = gen_compile_commands(f"{file_path}/..", cpps, soc_version)
    with open(f"{target_path}/compile_commands.json", "w") as fp:
        json.dump(commands, fp, indent="\t")
